#pragma once
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <map>

#include <boost/thread.hpp>
#include <boost/lambda/lambda.hpp>
#include <boost/lambda/bind.hpp>
#include <boost/lambda/construct.hpp>

template <class T, uint32_t N, uint32_t M, class init_functor > class MemoryPool;
template <class T, uint32_t N, uint32_t M, class init_functor >  class MemoryChunk;

typedef uint32_t POINTER_WIDTH; 

template <class T, uint32_t N, uint32_t M, class init_functor > class MemoryPool
{
	public:
		MemoryPool( init_functor init );
		~MemoryPool(void);

		inline uint64_t getUsedMem() const { return total_mem_used; }
		inline uint64_t getFreeMem() const { return total_mem_free; }
		inline uint64_t getUsedPackets() const { return total_blocks - free_blocks; }
		inline uint64_t getFreePackets() const { return free_blocks; }

		T *create();
		void destroy( T *t );
		void print() { printf( "Total Mem Used: %llubytes\nTotal Free Mem: %llubytes\nFree Blocks: %u\nTotal Blocks: %u\n" , total_mem_used, total_mem_free, free_blocks, total_blocks ); }
	private:
		void heapify();
		void heapify_up( uint32_t index );
		void allocate_chunks( init_functor init );
		MemoryChunk< T, N, M, init_functor > *chunks[N];
		std::map< POINTER_WIDTH, MemoryChunk< T, N, M, init_functor >*> pointer_map;
		uint64_t total_mem_free;
		uint64_t total_mem_used;
		uint32_t free_blocks;
		uint32_t total_blocks;
		boost::mutex thread_lock;
};

template <class T, uint32_t N, uint32_t M, class init_functor > MemoryPool< T, N, M, init_functor >::MemoryPool( init_functor init )
{
	total_mem_free = N * sizeof( T ) * ( M == 0 ? 1 : M );
	total_mem_used = 0;
	free_blocks = N;
	total_blocks = N;
	
	allocate_chunks( init );
}

template <class T, uint32_t N, uint32_t M, class init_functor > MemoryPool< T, N, M, init_functor >::~MemoryPool()
{
	thread_lock.lock();

		pointer_map.clear();

		for( int i = 0; i < N; i++ )
		{
			delete chunks[i];
			chunks[i] = 0;
		}

	thread_lock.unlock();
}

template <class T, uint32_t N, uint32_t M, class init_functor > void MemoryPool< T, N, M, init_functor >::allocate_chunks( init_functor init )
{
	for( uint32_t i = 0; i < N; i++ )
	{
		chunks[i] = new MemoryChunk< T, N, M, init_functor >( this, init, i );

		pointer_map.insert( std::pair< POINTER_WIDTH, MemoryChunk< T, N, M, init_functor >*>( (POINTER_WIDTH)chunks[i]->getPointer(), chunks[i] ) ); 
	}
}

template <class T, uint32_t N, uint32_t M, class init_functor > T *MemoryPool< T, N, M, init_functor >::create()
{
	T *ptr = 0;
	thread_lock.lock();
	
		if( !chunks[0]->allocated() )
		{
			ptr = chunks[0]->getPointer();

			chunks[0]->allocated( true );

			total_mem_used += sizeof( T ) * ( M == 0 ? 1 : M );
			total_mem_free -= sizeof( T ) * ( M == 0 ? 1 : M );
			free_blocks--;

			heapify();
		}

	thread_lock.unlock();

	return ptr;
}

template <class T, uint32_t N, uint32_t M, class init_functor > void MemoryPool< T, N, M, init_functor >::heapify()
{
	MemoryChunk< T, N, M, init_functor > *swap = 0;
	uint32_t index = 0;

	if( chunks[index]->allocated() )
	{
		while( index + 1 < total_blocks )
		{
			if( 2 * index + 1 < N && !chunks[2 * index + 1]->allocated() )
			{
				swap = chunks[2 * index + 1];
				
				uint32_t i_s = swap->getIndex();
				swap->setIndex( chunks[index]->getIndex() );
				chunks[index]->setIndex( i_s );

				chunks[2 * index + 1] = chunks[index];
				chunks[index] = swap;

				index = ( 2 * index + 1 );
			}
			else if( 2 * index + 2 < N && !chunks[2 * index + 2]->allocated() )
			{
				swap = chunks[2 * index + 2];

				uint32_t i_s = swap->getIndex();
				swap->setIndex( chunks[index]->getIndex() );
				chunks[index]->setIndex( i_s );

				chunks[2 * index + 2] = chunks[index];
				chunks[index] = swap;

				index = ( 2 * index + 2 );
			}
			else
				index = total_blocks;
		}
	}
}


template <class T, uint32_t N, uint32_t M, class init_functor > void MemoryPool< T, N, M, init_functor >::heapify_up( uint32_t index )
{
#define ABS( a ) ( (int32_t)(a) < 0 ? ( (int32_t)(a) * -1 ) : (a) )

	MemoryChunk< T, N, M, init_functor > *swap = 0;
	uint32_t i_s = 0;
	uint32_t i = index;

	if( !chunks[i]->allocated() )
	{
		while( i > 0 )
		{
			if( ABS( i - 1 ) / 2 >= 0 && chunks[ABS( i - 1 ) / 2]->allocated() )
			{
				swap = chunks[ABS( i - 1 ) / 2];

				i_s = swap->getIndex();
				swap->setIndex( chunks[i]->getIndex() );
				chunks[i]->setIndex( i_s );

				chunks[ABS( i - 1 ) / 2] = chunks[index];
				chunks[index] = swap;

				i = ABS( i - 1 ) / 2;
			}
			else if( ABS( i - 2 ) / 2 >= 0 && chunks[ABS( i - 2 ) / 2]->allocated() )
			{
				swap = chunks[ABS( i - 2 ) / 2];

				i_s = swap->getIndex();
				swap->setIndex( chunks[i]->getIndex() );
				chunks[i]->setIndex( i_s );

				chunks[ABS( i - 2 ) / 2] = chunks[index];
				chunks[index] = swap;

				i = ABS( i - 2 ) / 2;
			}
			else
				i = 0;
		}
	}
}

template <class T, uint32_t N, uint32_t M, class init_functor > void MemoryPool< T, N, M, init_functor >::destroy( T *t )
{
	thread_lock.lock();

		if( t != 0 )
		{
			pointer_map[(POINTER_WIDTH)t]->allocated( false );

			total_mem_used -= sizeof( T ) * ( M == 0 ? 1 : M );
			total_mem_free += sizeof( T ) * ( M == 0 ? 1 : M );
			free_blocks++;

			heapify_up( pointer_map[(POINTER_WIDTH)t]->getIndex() );
		}

	thread_lock.unlock();
}

template <class T, uint32_t N, uint32_t M, class init_functor > class MemoryChunk
{
	public:
		MemoryChunk( MemoryPool< T, N, M, init_functor > *pool, init_functor init, uint32_t _index );
		~MemoryChunk(void);

		T *getPointer() { return pointer; }
		
		uint64_t getSize() const { return size; }

		uint32_t getIndex() const { return index; }
		void setIndex( uint32_t _index ) { index = _index; }

		void destroy() { pool->destroy( this ); }
		bool allocated() const { return isAllocated; }
		void allocated( bool a ) { isAllocated = a; }

	private:
		T *pointer;
		MemoryPool< T, N, M, init_functor > *pool;
		bool isAllocated;
		bool isArray;
		uint64_t size;
		uint32_t index;
};

template <class T, uint32_t N, uint32_t M, class init_functor > MemoryChunk< T, N, M, init_functor >::MemoryChunk( MemoryPool< T, N, M, init_functor > *p, init_functor init, uint32_t _index ) : 
	pool( p ), index( _index ), isAllocated( false )
{
	size = sizeof( T );

	if( M  )
		isArray = true;
	else
		isArray = false;
	
	pointer = init();
}

template <class T, uint32_t N, uint32_t M, class init_functor > MemoryChunk< T, N, M, init_functor >::~MemoryChunk()
{
	if( M )
		delete [] pointer;
	else
		delete pointer;

	pointer = 0;
}
